""" 根据概率抽样出被选择的个体
原理:
1. 先归一化每个概率: prob[i]/sum 
2. 先随机一个种类i出来
3. 再随机一个数,看它和i的概率比大小,如果随机数更大,则拒绝

优势:
1. 避免了算累加"""
import numpy as np
import random
import collections

prob = [1, 2, 3, 4]
chr = ["a","b","c","d"]

def stochastic_accept(index_to_visit, transition_prob):
    """
    轮盘赌
    :param index_to_visit: a list of N index (list or tuple)
    :param transition_prob:
    :return: selected index
    """
    # calculate N and max fitness value
    N = len(index_to_visit)

    # normalize
    sum_tran_prob = np.sum(transition_prob)
    norm_transition_prob = transition_prob/sum_tran_prob

    # select: O(1)
    while True:
        # randomly select an individual with uniform probability
        # rand = random.random()
        # ind = int(N * rand)
        # if rand <= norm_transition_prob[ind]: # 转移概率 ind是可能被选中的路线
        ind = int(N * random.random())
        if random.random() < norm_transition_prob[ind]:
            return index_to_visit[ind]
    
if __name__ == "__main__":

    ans = []
    for i in range(100000):
        ans.append(stochastic_accept(index_to_visit=chr, transition_prob=prob))
    print(collections.Counter(ans))